{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "google",
   "metadata": {},
   "source": [
    "##### Copyright 2023 Google LLC."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "apache",
   "metadata": {},
   "source": [
    "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "you may not use this file except in compliance with the License.\n",
    "You may obtain a copy of the License at\n",
    "\n",
    "    http://www.apache.org/licenses/LICENSE-2.0\n",
    "\n",
    "Unless required by applicable law or agreed to in writing, software\n",
    "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "See the License for the specific language governing permissions and\n",
    "limitations under the License.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "basename",
   "metadata": {},
   "source": [
    "# school_scheduling_sat"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "link",
   "metadata": {},
   "source": [
    "<table align=\"left\">\n",
    "<td>\n",
    "<a href=\"https://colab.research.google.com/github/google/or-tools/blob/main/examples/notebook/contrib/school_scheduling_sat.ipynb\"><img src=\"https://raw.githubusercontent.com/google/or-tools/main/tools/colab_32px.png\"/>Run in Google Colab</a>\n",
    "</td>\n",
    "<td>\n",
    "<a href=\"https://github.com/google/or-tools/blob/main/examples/contrib/school_scheduling_sat.py\"><img src=\"https://raw.githubusercontent.com/google/or-tools/main/tools/github_32px.png\"/>View source on GitHub</a>\n",
    "</td>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "doc",
   "metadata": {},
   "source": [
    "First, you must install [ortools](https://pypi.org/project/ortools/) package in this colab."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "install",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install ortools"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "description",
   "metadata": {},
   "source": [
    "'''Solve a School Scheduling Problem'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "code",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ortools.sat.python import cp_model\n",
    "\n",
    "\n",
    "class SchoolSchedulingProblem():\n",
    "    '''Data of the problem.'''\n",
    "\n",
    "    def __init__(self, levels, sections, subjects, curriculum, teachers,\n",
    "                 specialties, time_slots):\n",
    "        self._levels = levels\n",
    "        self._sections = sections\n",
    "        self._subjects = subjects\n",
    "        self._curriculum = curriculum\n",
    "        assert len(self._curriculum) == len(self._levels) * len(\n",
    "            self._subjects), 'Some curriculum are missing'\n",
    "        for (lvl, sub) in self._curriculum.keys():\n",
    "            assert lvl in self._levels, f'{lvl} not in LEVELS'\n",
    "            assert sub in self._subjects, f'{sub} not in SUBJECTS'\n",
    "\n",
    "        self._teachers = teachers\n",
    "        self._specialties = specialties\n",
    "        assert len(self._specialties) == len(\n",
    "            self._subjects), 'Missing some rows'\n",
    "        for s, ts in self._specialties.items():\n",
    "            assert s in self._subjects, f'{s} is not in SUBJECTS'\n",
    "            for t in ts:\n",
    "                assert t in self._teachers, f'{t} is not in TEACHERS'\n",
    "\n",
    "        self._time_slots = time_slots\n",
    "\n",
    "    @property\n",
    "    def levels(self):\n",
    "        return self._levels\n",
    "\n",
    "    @property\n",
    "    def sections(self):\n",
    "        return self._sections\n",
    "\n",
    "    @property\n",
    "    def subjects(self):\n",
    "        return self._subjects\n",
    "\n",
    "    @property\n",
    "    def curriculum(self):\n",
    "        return self._curriculum\n",
    "\n",
    "    @property\n",
    "    def teachers(self):\n",
    "        return self._teachers\n",
    "\n",
    "    def teacher_name(self, teacher_idx):\n",
    "        assert 0 <= teacher_idx < len(self._teachers)\n",
    "        return list(self._teachers.keys())[teacher_idx]\n",
    "\n",
    "    def teacher_max_hours(self, teacher_idx):\n",
    "        assert 0 <= teacher_idx < len(self._teachers)\n",
    "        return list(self._teachers.values())[teacher_idx]\n",
    "\n",
    "    @property\n",
    "    def specialties(self):\n",
    "        return self._specialties\n",
    "\n",
    "    def specialtie_teachers(self, subject):\n",
    "        assert subject in self._subjects, f'{subject} not in SUBJECTS'\n",
    "        return self._specialties[subject]\n",
    "\n",
    "    @property\n",
    "    def time_slots(self):\n",
    "        return self._time_slots\n",
    "\n",
    "    def slot_duration(self, slot_idx):\n",
    "        assert 0 <= slot_idx < len(self._time_slots)\n",
    "        return list(self._time_slots.values())[slot_idx]\n",
    "\n",
    "\n",
    "class SchoolSchedulingSatSolver():\n",
    "    '''Solver instance.'''\n",
    "\n",
    "    def __init__(self, problem: SchoolSchedulingProblem):\n",
    "        # Problem\n",
    "        self._problem = problem\n",
    "\n",
    "        # Utilities\n",
    "        num_levels = len(self._problem.levels)\n",
    "        self._all_levels = range(num_levels)\n",
    "        num_sections = len(self._problem.sections)\n",
    "        self._all_sections = range(num_sections)\n",
    "        num_subjects = len(self._problem.subjects)\n",
    "        self._all_subjects = range(num_subjects)\n",
    "        num_teachers = len(self._problem.teachers)\n",
    "        self._all_teachers = range(num_teachers)\n",
    "        num_slots = len(self._problem.time_slots)\n",
    "        self._all_slots = range(num_slots)\n",
    "\n",
    "        # Create Model\n",
    "        self._model = cp_model.CpModel()\n",
    "\n",
    "        # Create Variables\n",
    "        self._assignment = {}\n",
    "        for lvl_idx, level in enumerate(self._problem.levels):\n",
    "            for sec_idx, section in enumerate(self._problem.sections):\n",
    "                for sub_idx, subject in enumerate(self._problem.subjects):\n",
    "                    for tch_idx, teacher in enumerate(self._problem.teachers):\n",
    "                        for slt_idx, slot in enumerate(self._problem.time_slots):\n",
    "                            key = (lvl_idx, sec_idx, sub_idx, tch_idx, slt_idx)\n",
    "                            name = f'{level}-{section} S:{subject} T:{teacher} Slot:{slot}'\n",
    "                            #print(name)\n",
    "                            if teacher in self._problem.specialtie_teachers(subject):\n",
    "                                self._assignment[key] = self._model.NewBoolVar(name)\n",
    "                            else:\n",
    "                                name = 'NO DISP ' + name\n",
    "                                self._assignment[key] = self._model.NewIntVar(0, 0, name)\n",
    "\n",
    "        # Constraints\n",
    "        # Each Level-Section must have the quantity of classes per Subject specified in the Curriculum\n",
    "        for lvl_idx, level in enumerate(self._problem.levels):\n",
    "            for sec_idx in self._all_sections:\n",
    "                for sub_idx, subject in enumerate(self._problem.subjects):\n",
    "                    required_duration = self._problem.curriculum[level, subject]\n",
    "                    #print(f'L:{level} S:{subject} duration:{required_duration}h')\n",
    "                    self._model.Add(\n",
    "                        sum(self._assignment[lvl_idx, sec_idx, sub_idx, tch_idx, slt_idx] *\n",
    "                            int(self._problem.slot_duration(slt_idx) * 10)\n",
    "                            for tch_idx in self._all_teachers\n",
    "                            for slt_idx in self._all_slots) == int(required_duration * 10))\n",
    "\n",
    "        # Each Level-Section can do at most one class at a time\n",
    "        for lvl_idx in self._all_levels:\n",
    "            for sec_idx in self._all_sections:\n",
    "                for slt_idx in self._all_slots:\n",
    "                    self._model.Add(\n",
    "                        sum([\n",
    "                            self._assignment[lvl_idx, sec_idx, sub_idx, tch_idx, slt_idx]\n",
    "                            for sub_idx in self._all_subjects\n",
    "                            for tch_idx in self._all_teachers\n",
    "                        ]) <= 1)\n",
    "\n",
    "        # Teacher can do at most one class at a time\n",
    "        for tch_idx in self._all_teachers:\n",
    "            for slt_idx in self._all_slots:\n",
    "                self._model.Add(\n",
    "                    sum([\n",
    "                        self._assignment[lvl_idx, sec_idx, sub_idx, tch_idx, slt_idx]\n",
    "                        for lvl_idx in self._all_levels\n",
    "                        for sec_idx in self._all_sections\n",
    "                        for sub_idx in self._all_subjects\n",
    "                    ]) <= 1)\n",
    "\n",
    "        # Maximum work hours for each teacher\n",
    "        for tch_idx in self._all_teachers:\n",
    "            self._model.Add(\n",
    "                sum([\n",
    "                    self._assignment[lvl_idx, sec_idx, sub_idx, tch_idx, slt_idx] *\n",
    "                    int(self._problem.slot_duration(slt_idx) * 10)\n",
    "                    for lvl_idx in self._all_levels\n",
    "                    for sec_idx in self._all_sections\n",
    "                    for sub_idx in self._all_subjects\n",
    "                    for slt_idx in self._all_slots\n",
    "                ]) <= int(self._problem.teacher_max_hours(tch_idx) * 10))\n",
    "\n",
    "        # Teacher makes all the classes of a subject's course\n",
    "        teacher_courses = {}\n",
    "        for lvl_idx, level in enumerate(self._problem.levels):\n",
    "            for sec_idx, section in enumerate(self._problem.sections):\n",
    "                for sub_idx, subject in enumerate(self._problem.subjects):\n",
    "                    for tch_idx, teacher in enumerate(self._problem.teachers):\n",
    "                        name = f'{level}-{section} S:{subject} T:{teacher}'\n",
    "                        teacher_courses[lvl_idx, sec_idx, sub_idx, tch_idx] = self._model.NewBoolVar(name)\n",
    "                        temp_array = [\n",
    "                            self._assignment[lvl_idx, sec_idx, sub_idx, tch_idx, slt_idx]\n",
    "                            for slt_idx in self._all_slots\n",
    "                        ]\n",
    "                        self._model.AddMaxEquality(\n",
    "                            teacher_courses[lvl_idx, sec_idx, sub_idx, tch_idx], temp_array)\n",
    "                    self._model.Add(\n",
    "                        sum([teacher_courses[lvl_idx, sec_idx, sub_idx, tch_idx]\n",
    "                             for tch_idx in self._all_teachers\n",
    "                        ]) == 1)\n",
    "\n",
    "\n",
    "    def print_teacher_schedule(self, tch_idx):\n",
    "        teacher_name = self._problem.teacher_name(tch_idx)\n",
    "        print(f'Teacher: {teacher_name}')\n",
    "        total_working_hours = 0\n",
    "        for slt_idx, slot in enumerate(self._problem.time_slots):\n",
    "            for lvl_idx, level in enumerate(self._problem.levels):\n",
    "                for sec_idx, section in enumerate(self._problem.sections):\n",
    "                    for sub_idx, subject in enumerate(self._problem.subjects):\n",
    "                        key = (lvl_idx, sec_idx, sub_idx, tch_idx, slt_idx)\n",
    "                        if self._solver.BooleanValue(self._assignment[key]):\n",
    "                            total_working_hours += self._problem.slot_duration(slt_idx)\n",
    "                            print(f'{slot}: C:{level}-{section} S:{subject}')\n",
    "        print(f'Total working hours: {total_working_hours}h')\n",
    "\n",
    "\n",
    "    def print_class_schedule(self, lvl_idx, sec_idx):\n",
    "        level = self._problem.levels[lvl_idx]\n",
    "        section = self._problem.sections[sec_idx]\n",
    "        print(f'Class: {level}-{section}')\n",
    "        total_working_hours = {}\n",
    "        for sub in self._problem.subjects:\n",
    "            total_working_hours[sub] = 0\n",
    "        for slt_idx, slot in enumerate(self._problem.time_slots):\n",
    "            for tch_idx, teacher in enumerate(self._problem.teachers):\n",
    "                for sub_idx, subject in enumerate(self._problem.subjects):\n",
    "                    key = (lvl_idx, sec_idx, sub_idx, tch_idx, slt_idx)\n",
    "                    if self._solver.BooleanValue(self._assignment[key]):\n",
    "                        total_working_hours[subject] += self._problem.slot_duration(slt_idx)\n",
    "                        print(f'{slot}: S:{subject} T:{teacher}')\n",
    "        for (subject, hours) in total_working_hours.items():\n",
    "            print(f'Total working hours for {subject}: {hours}h')\n",
    "\n",
    "\n",
    "    def print_school_schedule(self):\n",
    "        print('School:')\n",
    "        for slt_idx, slot in enumerate(self._problem.time_slots):\n",
    "            tmp = f'{slot}:'\n",
    "            for lvl_idx, level in enumerate(self._problem.levels):\n",
    "                for sec_idx, section in enumerate(self._problem.sections):\n",
    "                    for sub_idx, subject in enumerate(self._problem.subjects):\n",
    "                        for tch_idx, teacher in enumerate(self._problem.teachers):\n",
    "                            key = (lvl_idx, sec_idx, sub_idx, tch_idx, slt_idx)\n",
    "                            if self._solver.BooleanValue(self._assignment[key]):\n",
    "                                tmp += f' {level}-{section}:({subject},{teacher})'\n",
    "            print(tmp)\n",
    "\n",
    "\n",
    "    def solve(self):\n",
    "        print('Solving')\n",
    "        # Create Solver\n",
    "        self._solver = cp_model.CpSolver()\n",
    "\n",
    "        solution_printer = SchoolSchedulingSatSolutionPrinter()\n",
    "        status = self._solver.Solve(self._model, solution_printer)\n",
    "        print('Status: ', self._solver.StatusName(status))\n",
    "\n",
    "        if status == cp_model.OPTIMAL or status == cp_model.FEASIBLE:\n",
    "            print('\\n# Teachers')\n",
    "            for teacher_idx in self._all_teachers:\n",
    "                self.print_teacher_schedule(teacher_idx)\n",
    "\n",
    "            print('\\n# Classes')\n",
    "            for level_idx in self._all_levels:\n",
    "                for section_idx in self._all_sections:\n",
    "                    self.print_class_schedule(level_idx, section_idx)\n",
    "\n",
    "            print('\\n# School')\n",
    "            self.print_school_schedule()\n",
    "\n",
    "        print('Branches: ', self._solver.NumBranches())\n",
    "        print('Conflicts: ', self._solver.NumConflicts())\n",
    "        print('WallTime: ', self._solver.WallTime())\n",
    "\n",
    "\n",
    "class SchoolSchedulingSatSolutionPrinter(cp_model.CpSolverSolutionCallback):\n",
    "\n",
    "    def __init__(self):\n",
    "        cp_model.CpSolverSolutionCallback.__init__(self)\n",
    "        self.__solution_count = 0\n",
    "\n",
    "    def OnSolutionCallback(self):\n",
    "        print(\n",
    "            f'Solution #{self.__solution_count}, objective: {self.ObjectiveValue()}'\n",
    "        )\n",
    "        self.__solution_count += 1\n",
    "\n",
    "\n",
    "def main():\n",
    "    # DATA\n",
    "    ## Classes\n",
    "    LEVELS = [\n",
    "        '1',\n",
    "        '2',\n",
    "        '3',\n",
    "    ]\n",
    "    SECTIONS = [\n",
    "        'A',\n",
    "        'B',\n",
    "    ]\n",
    "    SUBJECTS = [\n",
    "        'English',\n",
    "        'Math',\n",
    "        #'Science',\n",
    "        'History',\n",
    "    ]\n",
    "    CURRICULUM = {\n",
    "        ('1', 'English'): 3,\n",
    "        ('1', 'Math'): 3,\n",
    "        ('1', 'History'): 2,\n",
    "        ('2', 'English'): 4,\n",
    "        ('2', 'Math'): 2,\n",
    "        ('2', 'History'): 2,\n",
    "        ('3', 'English'): 2,\n",
    "        ('3', 'Math'): 4,\n",
    "        ('3', 'History'): 2,\n",
    "    }\n",
    "\n",
    "    ## Teachers\n",
    "    TEACHERS = { # name, max_work_hours\n",
    "        'Mario': 14,\n",
    "        'Elvis': 12,\n",
    "        'Harry': 12,\n",
    "        'Ian': 14,\n",
    "    }\n",
    "    # Subject -> List of teachers who can teach it\n",
    "    SPECIALTIES = {\n",
    "        'English': ['Elvis', 'Ian'],\n",
    "        'Math': ['Mario', 'Ian'],\n",
    "        'History': ['Harry', 'Ian'],\n",
    "    }\n",
    "\n",
    "    ## Schedule\n",
    "    TIME_SLOTS = {\n",
    "            'Monday:08:00-09:30': 1.5,\n",
    "            'Monday:09:45-11:15': 1.5,\n",
    "            'Monday:11:30-12:30': 1,\n",
    "            'Monday:13:30-15:30': 2,\n",
    "            'Monday:15:45-17:15': 1.5,\n",
    "            'Tuesday:08:00-09:30': 1.5,\n",
    "            'Tuesday:09:45-11:15': 1.5,\n",
    "            'Tuesday:11:30-12:30': 1,\n",
    "            'Tuesday:13:30-15:30': 2,\n",
    "            'Tuesday:15:45-17:15': 1.5,\n",
    "            'Wednesday:08:00-09:30': 1.5,\n",
    "            'Wednesday:09:45-11:15': 1.5,\n",
    "            'Wednesday:11:30-12:30': 1,\n",
    "            'Thursday:08:00-09:30': 1.5,\n",
    "            'Thursday:09:45-11:15': 1.5,\n",
    "            'Thursday:11:30-12:30': 1,\n",
    "            'Thursday:13:30-15:30': 2,\n",
    "            'Thursday:15:45-17:15': 1.5,\n",
    "            'Friday:08:00-09:30': 1.5,\n",
    "            'Friday:09:45-11:15': 1.5,\n",
    "            'Friday:11:30-12:30': 1,\n",
    "            'Friday:13:30-15:30': 2,\n",
    "            'Friday:15:45-17:15': 1.5,\n",
    "            }\n",
    "\n",
    "    problem = SchoolSchedulingProblem(LEVELS, SECTIONS, SUBJECTS, CURRICULUM,\n",
    "                                      TEACHERS, SPECIALTIES, TIME_SLOTS)\n",
    "    solver = SchoolSchedulingSatSolver(problem)\n",
    "    solver.solve()\n",
    "\n",
    "\n",
    "main()\n",
    "\n"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}
